// TCP/UDP port forwarding

use std::net::SocketAddr;
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::{info, error};

pub struct TcpForwarder {
    listen_addr: SocketAddr,
    target_addr: String,
}

impl TcpForwarder {
    pub fn new(listen_addr: SocketAddr, target_addr: String) -> Self {
        Self {
            listen_addr,
            target_addr,
        }
    }

    pub async fn run(&self) -> anyhow::Result<()> {
        let listener = TcpListener::bind(self.listen_addr).await?;
        info!("TCP forwarding: {} -> {}", self.listen_addr, self.target_addr);

        loop {
            let (client, _) = listener.accept().await?;
            let target_addr = self.target_addr.clone();
            
            tokio::spawn(async move {
                if let Err(e) = Self::handle_connection(client, target_addr).await {
                    error!("TCP forward error: {}", e);
                }
            });
        }
    }

    async fn handle_connection(mut client: TcpStream, target_addr: String) -> anyhow::Result<()> {
        let mut target = TcpStream::connect(&target_addr).await?;

        let (mut client_read, mut client_write) = client.split();
        let (mut target_read, mut target_write) = target.split();

        let client_to_target = async {
            tokio::io::copy(&mut client_read, &mut target_write).await
        };

        let target_to_client = async {
            tokio::io::copy(&mut target_read, &mut client_write).await
        };

        tokio::try_join!(client_to_target, target_to_client)?;
        Ok(())
    }
}

pub struct UdpForwarder {
    listen_addr: SocketAddr,
    target_addr: String,
}

impl UdpForwarder {
    pub fn new(listen_addr: SocketAddr, target_addr: String) -> Self {
        Self {
            listen_addr,
            target_addr,
        }
    }

    pub async fn run(&self) -> anyhow::Result<()> {
        let socket = UdpSocket::bind(self.listen_addr).await?;
        info!("UDP forwarding: {} -> {}", self.listen_addr, self.target_addr);

        let target = UdpSocket::bind("0.0.0.0:0").await?;
        target.connect(&self.target_addr).await?;

        let mut buf = vec![0u8; 65535];
        
        loop {
            let (len, addr) = socket.recv_from(&mut buf).await?;
            let data = &buf[..len];
            
            // Forward to target
            target.send(data).await?;
            
            // Receive response
            let len = target.recv(&mut buf).await?;
            
            // Send back to client
            socket.send_to(&buf[..len], addr).await?;
        }
    }
}
